//! Trie data structure implementation in Rust.
//! This module provides a generic Trie that supports string and binary keys.

use std::collections::HashMap;
use std::fmt::Debug;
use std::hash::Hash;

/// A generic Trie data structure.
///
/// # Type Parameters
/// - `K`: The key type, which must implement `Eq`, `Hash`, and `Clone`.
/// - `V`: The value type, which must implement `Debug`.
pub struct Trie<K, V> {
    root: TrieNode<K, V>,
    entries: usize,
}

/// A node in the Trie.
struct TrieNode<K, V> {
    children: HashMap<K, TrieNode<K, V>>,
    value: Option<V>,
}

impl<K, V> Trie<K, V>
where
    K: Eq + Hash + Clone,
    V: Debug,
{
    /// Creates a new empty Trie.
    pub fn new() -> Self {
        Trie {
            root: TrieNode::new(),
            entries: 0,
        }
    }

    /// Inserts a key-value pair into the Trie.
    ///
    /// # Arguments
    /// - `key`: The key to insert.
    /// - `value`: The value to associate with the key.
    ///
    /// # Returns
    /// - `true` if the insertion was successful, `false` otherwise.
    pub fn insert(&mut self, key: Vec<K>, value: V) -> bool {
        let mut node = &mut self.root;
        for part in key {
            node = node.children.entry(part).or_insert_with(TrieNode::new);
        }
        if node.value.is_none() {
            self.entries += 1;
        }
        node.value = Some(value);
        true
    }

    /// Looks up a value by its key.
    ///
    /// # Arguments
    /// - `key`: The key to look up.
    ///
    /// # Returns
    /// - `Some(V)` if the key exists, `None` otherwise.
    pub fn lookup(&self, key: Vec<K>) -> Option<&V> {
        let mut node = &self.root;
        for part in key {
            if let Some(child) = node.children.get(&part) {
                node = child;
            } else {
                return None;
            }
        }
        node.value.as_ref()
    }

    /// Removes a key-value pair from the Trie.
    ///
    /// # Arguments
    /// - `key`: The key to remove.
    ///
    /// # Returns
    /// - `true` if the key was removed, `false` if it did not exist.
    pub fn remove(&mut self, key: Vec<K>) -> bool {
        let mut node = &mut self.root;
        let mut path = Vec::new();
        for part in key {
            path.push((node, part));
            if let Some(child) = node.children.get_mut(&part) {
                node = child;
            } else {
                return false;
            }
        }
        if node.value.is_some() {
            node.value = None;
            self.entries -= 1;
            for (parent, part) in path.into_iter().rev() {
                if parent.children[&part].children.is_empty() && parent.children[&part].value.is_none()
                {
                    parent.children.remove(&part);
                } else {
                    break;
                }
            }
            true
        } else {
            false
        }
    }

    /// Returns the number of entries in the Trie.
    pub fn num_entries(&self) -> usize {
        self.entries
    }
}

impl<K, V> TrieNode<K, V> {
    /// Creates a new empty TrieNode.
    fn new() -> Self {
        TrieNode {
            children: HashMap::new(),
            value: None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const NUM_TEST_VALUES: usize = 10000;

    #[test]
    fn test_trie_new_free() {
        let mut trie: Trie<char, &str> = Trie::new();

        // Allocate and free an empty trie
        assert!(trie.num_entries() == 0);

        // Add some values before freeing
        assert!(trie.insert("hello".chars().collect(), "there"));
        assert!(trie.insert("hell".chars().collect(), "testing"));
        assert!(trie.insert("testing".chars().collect(), "testing"));
        assert!(trie.insert("".chars().collect(), "asfasf"));
        assert_eq!(trie.num_entries(), 4);

        // Add a value, remove it and then free
        assert!(trie.insert("hello".chars().collect(), "there"));
        assert!(trie.remove("hello".chars().collect()));
        assert_eq!(trie.num_entries(), 3);
    }

    #[test]
    fn test_trie_insert() {
        let mut trie = generate_trie();

        // Test insert of NULL value has no effect
        let entries = trie.num_entries();
        assert!(!trie.insert("hello world".chars().collect(), ""));
        assert_eq!(trie.num_entries(), entries);
    }

    #[test]
    fn test_trie_lookup() {
        let trie = generate_trie();

        // Test lookup for non-existent values
        assert_eq!(trie.lookup("000000000000000".chars().collect()), None);
        assert_eq!(trie.lookup("".chars().collect()), None);

        // Look up all values
        for i in 0..NUM_TEST_VALUES {
            let key: String = format!("{}", i);
            let value = trie.lookup(key.chars().collect());
            assert_eq!(value, Some(&i));
        }
    }

    #[test]
    fn test_trie_remove() {
        let mut trie = generate_trie();

        // Test remove on non-existent values
        assert!(!trie.remove("000000000000000".chars().collect()));
        assert!(!trie.remove("".chars().collect()));

        let entries = trie.num_entries();
        assert_eq!(entries, NUM_TEST_VALUES);

        // Remove all values
        for i in 0..NUM_TEST_VALUES {
            let key: String = format!("{}", i);
            assert!(trie.remove(key.chars().collect()));
            assert_eq!(trie.num_entries(), entries - i - 1);
        }
    }

    #[test]
    fn test_trie_replace() {
        let mut trie = generate_trie();

        // Test replacing values
        let val = 999;
        assert!(trie.insert("999".chars().collect(), val));
        assert_eq!(trie.num_entries(), NUM_TEST_VALUES);
        assert_eq!(trie.lookup("999".chars().collect()), Some(&val));
    }

    #[test]
    fn test_trie_insert_empty() {
        let mut trie: Trie<char, &str> = Trie::new();

        // Test insert on empty string
        assert!(trie.insert("".chars().collect(), "buf"));
        assert_ne!(trie.num_entries(), 0);
        assert_eq!(trie.lookup("".chars().collect()), Some(&"buf"));
        assert!(trie.remove("".chars().collect()));
        assert_eq!(trie.num_entries(), 0);
    }

    #[test]
    fn test_trie_free_long() {
        let long_string = vec!['A'; 4096];
        let mut trie: Trie<char, Vec<char>> = Trie::new();

        // Create a trie and add the string
        trie.insert(long_string.clone(), long_string.clone());
        assert_eq!(trie.num_entries(), 1);
    }

    #[test]
    fn test_trie_negative_keys() {
        let mut trie: Trie<char, &str> = Trie::new();
        let my_key: Vec<char> = vec!['a', 'b', 'c', 206 as char, 236 as char, '\0'];

        assert!(trie.insert(my_key.clone(), "hello world"));
        assert_eq!(trie.lookup(my_key.clone()), Some(&"hello world"));
        assert!(trie.remove(my_key.clone()));
        assert!(!trie.remove(my_key.clone()));
        assert_eq!(trie.lookup(my_key.clone()), None);
    }

    fn generate_trie() -> Trie<char, i32> {
        let mut trie = Trie::new();
        for i in 0..NUM_TEST_VALUES {
            let key: String = format!("{}", i);
            trie.insert(key.chars().collect(), i);
        }
        trie
    }
}
